package cn.maihe.elg.operation.utils;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.macs.CBCBlockCipherMac;
import org.bouncycastle.crypto.macs.GMac;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.paddings.BlockCipherPadding;
import org.bouncycastle.crypto.paddings.PKCS7Padding;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;

/**
 * SM4国密算法工具类，提供加密、解密、MAC计算等功能。
 * 支持ECB和CBC模式，支持带填充和不带填充的处理方式。
 *
 * @Description SM4国密算法工具类
 * @Author WJH
 * @Date 2021/07/12
 */
public class SM4Util {

    // 添加BouncyCastle安全提供者
    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static final String ALGORITHM_NAME = "SM4";
    public static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
    public static final String ALGORITHM_NAME_ECB_NOPADDING = "SM4/ECB/NoPadding";
    public static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS5Padding";
    public static final String ALGORITHM_NAME_CBC_NOPADDING = "SM4/CBC/NoPadding";

    /**
     * SM4算法目前只支持128位（即密钥16字节）
     */
    public static final int DEFAULT_KEY_SIZE = 128;

    /**
     * 生成默认长度的SM4密钥
     *
     * @return 生成的密钥字节数组
     * @throws NoSuchAlgorithmException 当算法不支持时抛出
     * @throws NoSuchProviderException 当安全提供者不支持时抛出
     */
    public static byte[] generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
        return generateKey(DEFAULT_KEY_SIZE);
    }

    /**
     * 生成指定长度的SM4密钥
     *
     * @param keySize 密钥长度（位），SM4算法推荐使用128位
     * @return 生成的密钥字节数组
     * @throws NoSuchAlgorithmException 当算法不支持时抛出
     * @throws NoSuchProviderException 当安全提供者不支持时抛出
     */
    public static byte[] generateKey(int keySize) throws NoSuchAlgorithmException, NoSuchProviderException {
        KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM_NAME, BouncyCastleProvider.PROVIDER_NAME);
        kg.init(keySize, new SecureRandom());
        return kg.generateKey().getEncoded();
    }

    /**
     * 使用ECB模式和PKCS5Padding填充进行SM4加密
     *
     * @param key 待加密数据的密钥
     * @param data 待加密的数据
     * @return 加密后的数据
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     */
    public static byte[] encrypt_ECB_Padding(byte[] key, byte[] data)
            throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException,
            NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
        Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.ENCRYPT_MODE, key);
        return cipher.doFinal(data);
    }

    /**
     * 使用ECB模式和PKCS5Padding填充进行SM4解密
     *
     * @param key 解密密钥
     * @param cipherText 待解密的数据
     * @return 解密后的明文数据
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     */
    public static byte[] decrypt_ECB_Padding(byte[] key, byte[] cipherText)
            throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
            NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
        Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(cipherText);
    }

    /**
     * 使用ECB模式和无填充进行SM4加密
     *
     * @param key 加密密钥
     * @param data 待加密数据（长度必须为16字节的整数倍）
     * @return 加密后的数据
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     */
    public static byte[] encrypt_ECB_NoPadding(byte[] key, byte[] data)
            throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException,
            NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
        Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_NOPADDING, Cipher.ENCRYPT_MODE, key);
        return cipher.doFinal(data);
    }

    /**
     * 使用ECB模式和无填充进行SM4解密
     *
     * @param key 解密密钥
     * @param cipherText 待解密数据（长度必须为16字节的整数倍）
     * @return 解密后的明文数据
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     */
    public static byte[] decrypt_ECB_NoPadding(byte[] key, byte[] cipherText)
            throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
            NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
        Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_NOPADDING, Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(cipherText);
    }

    /**
     * 使用CBC模式和PKCS5Padding填充进行SM4加密
     *
     * @param key 加密密钥
     * @param iv 初始化向量
     * @param data 待加密数据
     * @return 加密后的数据
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidAlgorithmParameterException 算法参数无效时抛出
     */
    public static byte[] encrypt_CBC_Padding(byte[] key, byte[] iv, byte[] data)
            throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException,
            NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException,
            InvalidAlgorithmParameterException {
        Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.ENCRYPT_MODE, key, iv);
        return cipher.doFinal(data);
    }

    /**
     * 使用CBC模式和PKCS5Padding填充进行SM4解密
     *
     * @param key 解密密钥
     * @param iv 初始化向量
     * @param cipherText 待解密数据
     * @return 解密后的明文数据
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws InvalidAlgorithmParameterException 算法参数无效时抛出
     */
    public static byte[] decrypt_CBC_Padding(byte[] key, byte[] iv, byte[] cipherText)
            throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
            NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException,
            InvalidAlgorithmParameterException {
        Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.DECRYPT_MODE, key, iv);
        return cipher.doFinal(cipherText);
    }

    /**
     * 使用CBC模式和无填充进行SM4加密
     *
     * @param key 加密密钥
     * @param iv 初始化向量
     * @param data 待加密数据（长度必须为16字节的整数倍）
     * @return 加密后的数据
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidAlgorithmParameterException 算法参数无效时抛出
     */
    public static byte[] encrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] data)
            throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException,
            NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException,
            InvalidAlgorithmParameterException {
        Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_NOPADDING, Cipher.ENCRYPT_MODE, key, iv);
        return cipher.doFinal(data);
    }

    /**
     * 使用CBC模式和无填充进行SM4解密
     *
     * @param key 解密密钥
     * @param iv 初始化向量
     * @param cipherText 待解密数据（长度必须为16字节的整数倍）
     * @return 解密后的明文数据
     * @throws IllegalBlockSizeException 块大小非法时抛出
     * @throws BadPaddingException 填充错误时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws InvalidAlgorithmParameterException 算法参数无效时抛出
     */
    public static byte[] decrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] cipherText)
            throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
            NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException,
            InvalidAlgorithmParameterException {
        Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_NOPADDING, Cipher.DECRYPT_MODE, key, iv);
        return cipher.doFinal(cipherText);
    }

    /**
     * 计算SM4-CMAC消息认证码
     *
     * @param key MAC计算密钥
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     */
    public static byte[] doCMac(byte[] key, byte[] data) throws NoSuchProviderException, NoSuchAlgorithmException,
            InvalidKeyException {
        Key keyObj = new SecretKeySpec(key, ALGORITHM_NAME);
        return doMac("SM4-CMAC", keyObj, data);
    }

    /**
     * 计算SM4-GMAC消息认证码
     *
     * @param key MAC计算密钥
     * @param iv 初始化向量
     * @param tagLength 标签长度（字节）
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     */
    public static byte[] doGMac(byte[] key, byte[] iv, int tagLength, byte[] data) {
        org.bouncycastle.crypto.Mac mac = new GMac(new GCMBlockCipher(new SM4Engine()), tagLength * 8);
        return doMac(mac, key, iv, data);
    }

    /**
     * 默认使用PKCS7Padding/PKCS5Padding填充的CBCMAC
     *
     * @param key MAC计算密钥
     * @param iv 初始化向量
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     */
    public static byte[] doCBCMac(byte[] key, byte[] iv, byte[] data) {
        SM4Engine engine = new SM4Engine();
        org.bouncycastle.crypto.Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, new PKCS7Padding());
        return doMac(mac, key, iv, data);
    }

    /**
     * 使用指定填充方式计算CBCMAC
     *
     * @param key MAC计算密钥
     * @param iv 初始化向量
     * @param padding 填充方式，可以传null表示NoPadding，此时调用方需保证数据长度为BlockSize的整数倍
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     * @throws Exception 当填充为null且数据长度不是BlockSize整数倍时抛出异常
     */
    public static byte[] doCBCMac(byte[] key, byte[] iv, BlockCipherPadding padding, byte[] data) throws Exception {
        SM4Engine engine = new SM4Engine();
        if (padding == null) {
            if (data.length % engine.getBlockSize() != 0) {
                throw new Exception("if no padding, data length must be multiple of SM4 BlockSize");
            }
        }
        org.bouncycastle.crypto.Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, padding);
        return doMac(mac, key, iv, data);
    }


    /**
     * 使用BouncyCastle的MAC接口计算消息认证码
     *
     * @param mac MAC计算接口
     * @param key MAC计算密钥
     * @param iv 初始化向量
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     */
    private static byte[] doMac(org.bouncycastle.crypto.Mac mac, byte[] key, byte[] iv, byte[] data) {
        CipherParameters cipherParameters = new KeyParameter(key);
        mac.init(new ParametersWithIV(cipherParameters, iv));
        mac.update(data, 0, data.length);
        byte[] result = new byte[mac.getMacSize()];
        mac.doFinal(result, 0);
        return result;
    }

    /**
     * 使用Java标准API计算消息认证码
     *
     * @param algorithmName 算法名称
     * @param key MAC计算密钥
     * @param data 待计算MAC的数据
     * @return 计算得到的MAC值
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     */
    private static byte[] doMac(String algorithmName, Key key, byte[] data) throws NoSuchProviderException,
            NoSuchAlgorithmException, InvalidKeyException {
        Mac mac = Mac.getInstance(algorithmName, BouncyCastleProvider.PROVIDER_NAME);
        mac.init(key);
        mac.update(data);
        return mac.doFinal();
    }

    /**
     * 生成ECB模式的Cipher对象
     *
     * @param algorithmName 算法名称
     * @param mode 加密/解密模式
     * @param key 密钥
     * @return 生成的Cipher对象
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     * @throws InvalidKeyException 密钥无效时抛出
     */
    private static Cipher generateECBCipher(String algorithmName, int mode, byte[] key)
            throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException,
            InvalidKeyException {
        Cipher cipher = Cipher.getInstance(algorithmName, BouncyCastleProvider.PROVIDER_NAME);
        Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        cipher.init(mode, sm4Key);
        return cipher;
    }

    /**
     * 生成CBC模式的Cipher对象
     *
     * @param algorithmName 算法名称
     * @param mode 加密/解密模式
     * @param key 密钥
     * @param iv 初始化向量
     * @return 生成的Cipher对象
     * @throws InvalidKeyException 密钥无效时抛出
     * @throws InvalidAlgorithmParameterException 算法参数无效时抛出
     * @throws NoSuchAlgorithmException 算法不支持时抛出
     * @throws NoSuchProviderException 安全提供者不支持时抛出
     * @throws NoSuchPaddingException 填充方式不支持时抛出
     */
    private static Cipher generateCBCCipher(String algorithmName, int mode, byte[] key, byte[] iv)
            throws InvalidKeyException, InvalidAlgorithmParameterException, NoSuchAlgorithmException,
            NoSuchProviderException, NoSuchPaddingException {
        Cipher cipher = Cipher.getInstance(algorithmName, BouncyCastleProvider.PROVIDER_NAME);
        Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(mode, sm4Key, ivParameterSpec);
        return cipher;
    }
}
